﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.ServiceProcess;
using System.Timers;

namespace VDIMonitorService
{

    #region User Session Enums

    [StructLayout(LayoutKind.Sequential)]
    internal struct WTS_SESSION_INFO
    {
        public Int32 SessionID;
        [MarshalAs(UnmanagedType.LPStr)] public String pWinStationName;
        public WTS_CONNECTSTATE_CLASS State;
    }

    internal enum WTS_CONNECTSTATE_CLASS
    {
        WTSActive,
        WTSConnected,
        WTSConnectQuery,
        WTSShadow,
        WTSDisconnected,
        WTSIdle,
        WTSListen,
        WTSReset,
        WTSDown,
        WTSInit
    }

    internal enum WTS_INFO_CLASS
    {
        WTSInitialProgram,
        WTSApplicationName,
        WTSWorkingDirectory,
        WTSOEMId,
        WTSSessionId,
        WTSUserName,
        WTSWinStationName,
        WTSDomainName,
        WTSConnectState,
        WTSClientBuildNumber,
        WTSClientName,
        WTSClientDirectory,
        WTSClientProductId,
        WTSClientHardwareId,
        WTSClientAddress,
        WTSClientDisplay,
        WTSClientProtocolType,
        WTSIdleTime,
        WTSLogonTime,
        WTSIncomingBytes,
        WTSOutgoingBytes,
        WTSIncomingFrames,
        WTSOutgoingFrames,
        WTSClientInfo,
        WTSSessionInfo
    }

    #endregion

    /// <summary>
    ///     This services will monitor all active sessions for an instance of the monitored process name. If no instances are found, the user is logged off
    /// </summary>
    public partial class VdiMonitorService : ServiceBase
    {
        private const int TimerIntervalInMilliseconds = 1000;
        private const string MonitoredProcessName = "WSWC";
        private readonly Timer ProcessPollTimer = new Timer();
        private IntPtr LocalSessionServer;
        private bool ProcessMonitored;

        public VdiMonitorService()
        {
            InitializeComponent();
        }

        protected override void OnStart(string[] args)
        {
            try
            {
                LocalSessionServer = WTSOpenServer(Environment.MachineName);
                ProcessMonitored = false;
                ProcessPollTimer.Elapsed += ProcessPollTimerElapsed;
                //Interval is in milliseconds 
                ProcessPollTimer.Interval = TimerIntervalInMilliseconds;
                ProcessPollTimer.Start();
            }
            catch (Exception e)
            {
                Logger.LogException(e);
            }
        }

        private void ProcessPollTimerElapsed(object sender, ElapsedEventArgs e)
        {            
            try
            {                
                if (ActiveUserLogin(LocalSessionServer))
                {
                    RegisterProcessExitHandler();
                }
            }
            catch (Exception exception)
            {
                Logger.LogException(exception);
            }
        }

        private void RegisterProcessExitHandler()
        {
            if (ProcessRunning(MonitoredProcessName) && !ProcessMonitored)
            {
                Process p = Process.GetProcessesByName(MonitoredProcessName).First();
                p.EnableRaisingEvents = true;
                p.Exited += p_Exited;
                ProcessMonitored = true;
                ProcessPollTimer.Stop();
                Logger.LogInfo("Monitoring process for exit. Stopped polling Timer.");
            }
        }

        void p_Exited(object sender, EventArgs e)
        {
            ProcessMonitored = false;
            LogOffUserSession(LocalSessionServer);
            ProcessPollTimer.Start();
            Logger.LogInfo("Restarting polling");
        }

        private static bool ProcessRunning(string procName)
        {
            return Process.GetProcessesByName(procName).Any();
        }

        private static bool ActiveUserLogin(IntPtr server)
        {
            Dictionary<string, int> userSessionDict = GetUserSessionDictionary(server, GetSessionIDs(server));
            return userSessionDict.Any();
        }

        private static bool LogOffUserSession(IntPtr server)
        {
            try
            {
                Logger.RecordUserLogoff("Initiating user logoff.");
                List<int> sessions = GetSessionIDs(server);
                Dictionary<string, int> userSessionDictionary = GetUserSessionDictionary(server, sessions);
                return WTSLogoffSession(server, userSessionDictionary.First().Value, true);
            }
            catch (Exception e)
            {
                Logger.LogException(e);
            }

            return false;
        }

        protected override void OnStop()
        {
            WTSCloseServer(LocalSessionServer);
        }

        #region User Session Code

        [DllImport("wtsapi32.dll", SetLastError = true)]
        private static extern bool WTSLogoffSession(IntPtr hServer, int SessionId, bool bWait);

        [DllImport("Wtsapi32.dll")]
        private static extern bool WTSQuerySessionInformation(
            IntPtr hServer, int sessionId, WTS_INFO_CLASS wtsInfoClass, out IntPtr ppBuffer, out uint pBytesReturned);

        [DllImport("wtsapi32.dll", SetLastError = true)]
        private static extern IntPtr WTSOpenServer([MarshalAs(UnmanagedType.LPStr)] String pServerName);

        [DllImport("wtsapi32.dll")]
        private static extern void WTSCloseServer(IntPtr hServer);

        [DllImport("wtsapi32.dll", SetLastError = true)]
        private static extern Int32 WTSEnumerateSessions(IntPtr hServer, [MarshalAs(UnmanagedType.U4)] Int32 Reserved,
                                                         [MarshalAs(UnmanagedType.U4)] Int32 Version,
                                                         ref IntPtr ppSessionInfo,
                                                         [MarshalAs(UnmanagedType.U4)] ref Int32 pCount);

        [DllImport("wtsapi32.dll")]
        private static extern void WTSFreeMemory(IntPtr pMemory);

        internal static List<int> GetSessionIDs(IntPtr server)
        {
            var sessionIds = new List<int>();
            IntPtr buffer = IntPtr.Zero;
            int count = 0;
            int retval = WTSEnumerateSessions(server, 0, 1, ref buffer, ref count);
            int dataSize = Marshal.SizeOf(typeof (WTS_SESSION_INFO));
            Int64 current = (int) buffer;

            if (retval != 0)
            {
                for (int i = 0; i < count; i++)
                {
                    var si = (WTS_SESSION_INFO) Marshal.PtrToStructure((IntPtr) current, typeof (WTS_SESSION_INFO));
                    current += dataSize;
                    sessionIds.Add(si.SessionID);
                }
                WTSFreeMemory(buffer);
            }
            return sessionIds;
        }

        private static Dictionary<string, int> GetUserSessionDictionary(IntPtr server, List<int> sessions)
        {
            var userSession = new Dictionary<string, int>();

            foreach (int sessionId in sessions)
            {
                string uName = GetUserName(sessionId, server);
                if (!string.IsNullOrEmpty(uName))
                    userSession.Add(uName, sessionId);
            }
            return userSession;
        }

        internal static string GetUserName(int sessionId, IntPtr server)
        {
            IntPtr buffer = IntPtr.Zero;
            uint count = 0;
            string userName = string.Empty;
            try
            {
                WTSQuerySessionInformation(server, sessionId, WTS_INFO_CLASS.WTSUserName, out buffer, out count);
                userName = Marshal.PtrToStringAnsi(buffer).ToUpper().Trim();
            }
            finally
            {
                WTSFreeMemory(buffer);
            }
            return userName;
        }

        #endregion
    }
}